#!/usr/bin/env python3

from __future__ import print_function

import argparse
import base64
import copy
import hashlib
import json
import logging
import os
import pickle
import pkgutil
import signal
import sqlite3
import ssl
import subprocess
import sys
import time
from datetime import datetime, timezone
from time import sleep

from flask import Flask, request, jsonify, make_response, abort, url_for, g
from flask.json import JSONEncoder
from flask_socketio import SocketIO, emit

# Empire imports
from lib.common import empire, helpers, users
from lib.common.empire import MainMenu

# Check if running Python 3
if sys.version[0] == '2':
    print(helpers.color("[!] Please use Python 3"))
    sys.exit()

global serverExitCommand
serverExitCommand = 'restart'


#####################################################
#
# Database interaction methods for the RESTful API
#
#####################################################


def database_check_docker():
    """
    Check for docker and setup database if nessary.
    """
    if os.path.exists('/.dockerenv'):
        if not os.path.exists('data/empire.db'):
            print('[*] Fresh start in docker, running reset.sh for you')
            subprocess.call(['./setup/reset.sh'])


def adapt_datetime(val):
    """
    This adapter is available in sqlite3/dbapi2.py, but uses a " " as the seperator
    This uses the default of "T" which fits ISO-8601
    """
    return val.isoformat()


def convert_timestamp(val):
    """
    The original version of this in sqlite3/dbapi2.py doesn't account for timezone aware datetimes. Using fromisoformat
    should handle both naive and aware and astimezone will force naive timestamps to utc.
    datetimes.
    """
    return datetime.fromisoformat(val.decode('utf-8')).astimezone(timezone.utc)


class MyJsonEncoder(JSONEncoder):
    def default(self, o):
        if isinstance(o, datetime):
            return o.isoformat()

        return super().default(o)


def database_connect():
    """
    Connect with the backend ./empire.db sqlite database and return the
    connection object.
    """
    try:
        sqlite3.register_adapter(datetime, adapt_datetime)
        sqlite3.register_converter("timestamp", convert_timestamp)
        # set the database connectiont to autocommit w/ isolation level
        conn = sqlite3.connect('./data/empire.db', check_same_thread=False, detect_types=sqlite3.PARSE_DECLTYPES)
        conn.text_factory = str
        conn.isolation_level = None
        return conn

    except Exception:
        print(helpers.color("[!] Could not connect to database"))
        print(helpers.color("[!] Please run setup_database.py"))
        sys.exit()


def execute_db_query(conn, query, args=None):
    """
    Execute the supplied query on the provided db conn object
    with optional args for a paramaterized query.
    """
    cur = conn.cursor()
    if args:
        cur.execute(query, args)
    else:
        cur.execute(query)
    results = cur.fetchall()
    cur.close()
    return results

####################################################################
#
# The Empire RESTful API. To see more information about it, check out the official wiki.
#
# Adapted from http://blog.miguelgrinberg.com/post/designing-a-restful-api-with-python-and-flask.
# Example code at https://gist.github.com/miguelgrinberg/5614326.
#
#
#    Verb     URI                                                   Action
#    ----     ---                                                   ------
#    GET      http://localhost:1337/api/version                     return the current Empire version
#    GET      http://localhost:1337/api/map                         return list of all API routes
#    GET      http://localhost:1337/api/config                      return the current default config
#
#    GET      http://localhost:1337/api/stagers                     return all current stagers
#    GET      http://localhost:1337/api/stagers/X                   return the stager with name X
#    POST     http://localhost:1337/api/stagers                     generate a stager given supplied options (need to implement)
#
#    GET      http://localhost:1337/api/modules                     return all current modules
#    GET      http://localhost:1337/api/modules/<name>              return the module with the specified name
#    POST     http://localhost:1337/api/modules/<name>              execute the given module with the specified options
#    POST     http://localhost:1337/api/modules/search              searches modulesfor a passed term
#    POST     http://localhost:1337/api/modules/search/modulename   searches module names for a specific term
#    POST     http://localhost:1337/api/modules/search/description  searches module descriptions for a specific term
#    POST     http://localhost:1337/api/modules/search/comments     searches module comments for a specific term
#    POST     http://localhost:1337/api/modules/search/author       searches module authors for a specific term
#
#    GET      http://localhost:1337/api/listeners                   return all current listeners
#    GET      http://localhost:1337/api/listeners/Y                 return the listener with id Y
#    DELETE   http://localhost:1337/api/listeners/Y                 kills listener Y
#    GET      http://localhost:1337/api/listeners/types             returns a list of the loaded listeners that are available for use
#    GET      http://localhost:1337/api/listeners/options/Y         return listener options for Y
#    POST     http://localhost:1337/api/listeners/Y                 starts a new listener with the specified options
#
#    GET      http://localhost:1337/api/agents                      return all current agents
#    GET      http://localhost:1337/api/agents/stale                return all stale agents
#    DELETE   http://localhost:1337/api/agents/stale                removes stale agents from the database
#    DELETE   http://localhost:1337/api/agents/Y                    removes agent Y from the database
#    GET      http://localhost:1337/api/agents/Y                    return the agent with name Y
#    GET      http://localhost:1337/api/agents/Y/directory          return the directory with the name given by the query parameter 'directory'
#    POST     http://localhost:1337/api/agents/Y/directory          task the agent Y to scrape the directory given by the query parameter 'directory'
#    GET      http://localhost:1337/api/agents/Y/results            return tasking results for the agent with name Y
#    DELETE   http://localhost:1337/api/agents/Y/results            deletes the result buffer for agent Y
#    GET      http://localhost:1337/api/agents/Y/task/Z             return the tasking Z for agent Y
#    POST     http://localhost:1337/api/agents/Y/download           task agent Y to download a file
#    POST     http://localhost:1337/api/agents/Y/upload             task agent Y to upload a file
#    POST     http://localhost:1337/api/agents/Y/shell              task agent Y to execute a shell command
#    POST     http://localhost:1337/api/agents/Y/rename             rename agent Y
#    GET/POST http://localhost:1337/api/agents/Y/clear              clears the result buffer for agent Y
#    GET/POST http://localhost:1337/api/agents/Y/kill               kill agent Y
#
#    GET      http://localhost:1337/api/creds                       return stored credentials
#    POST     http://localhost:1337/api/creds                       add creds to the database
#
#    GET      http://localhost:1337/api/reporting                   return all logged events
#    GET      http://localhost:1337/api/reporting/agent/X           return all logged events for the given agent name X
#    GET      http://localhost:1337/api/reporting/type/Y            return all logged events of type Y (checkin, task, result, rename)
#    GET      http://localhost:1337/api/reporting/msg/Z             return all logged events matching message Z, wildcards accepted
#
#
#    POST     http://localhost:1337/api/admin/login                 retrieve the API token given the correct username and password
#    POST     http://localhost:1337/api/admin/logout                logout of current user account
#    GET      http://localhost:1337/api/admin/restart               restart the RESTful API
#    GET      http://localhost:1337/api/admin/shutdown              shutdown the RESTful API
#
#    GET      http://localhost:1337/api/users                       return all users from database
#    GET      http://localhost:1337/api/users/X                     return the user with id X
#    GET      http://localhost:1337/api/users/me                    return the user for the given token
#    POST     http://localhost:1337/api/users                       add a new user
#    PUT      http://localhost:1337/api/users/Y/disable             disable/enable user Y
#    PUT      http://localhost:1337/api/users/Y/updatepassword      update password for user Y
#
####################################################################

def start_restful_api(empireMenu: MainMenu, suppress=False, username=None, password=None, port=1337):
    """
    Kick off the RESTful API with the given parameters.

    empireMenu  -   Main empire menu object
    suppress    -   suppress most console output
    username    -   optional username to use for the API, otherwise pulls from the empire.db config
    password    -   optional password to use for the API, otherwise pulls from the empire.db config
    port        -   port to start the API on, defaults to 1337 ;)
    """
    app = Flask(__name__)

    app.json_encoder = MyJsonEncoder

    conn = database_connect()

    main = empireMenu

    global serverExitCommand

    # if a username/password were not supplied, use the creds stored in the db
    #(dbUsername, dbPassword) = execute_db_query(conn, "SELECT api_username, api_password FROM config")[0]

    if username:
        main.users.update_username(1, username[0])

    if password:
        main.users.update_password(1, password[0])

    print('')
    print(" * Starting Empire RESTful API on port: %s" % (port))

    oldStdout = sys.stdout
    if suppress:
        # suppress the normal Flask output
        log = logging.getLogger('werkzeug')
        log.setLevel(logging.ERROR)

        # suppress all stdout and don't initiate the main cmdloop
        sys.stdout = open(os.devnull, 'w')

    # validate API token before every request except for the login URI
    @app.before_request
    def check_token():
        """
        Before every request, check if a valid token is passed along with the request.
        """
        try:
            if request.path != '/api/admin/login':
                token = request.args.get('token')
                if token and len(token) > 0:
                    user = main.users.get_user_from_token(token)
                    if user:
                        g.user = user
                    else:
                        return make_response('', 401)
                else:
                    return make_response('', 401)
        except:
            return make_response('', 401)

    @app.after_request
    def add_cors(response):
        response.headers['Access-Control-Allow-Origin'] = '*'
        return response

    @app.errorhandler(Exception)
    def exception_handler(error):
        """
        Generic exception handler.
        """
        code = error.code if hasattr(error, 'code') else '500'
        return make_response(jsonify({'error': repr(error)}), code)

    @app.errorhandler(404)
    def not_found(error):
        """
        404/not found handler.
        """
        return make_response(jsonify({'error': 'Not found'}), 404)

    @app.route('/api/version', methods=['GET'])
    def get_version():
        """
        Returns the current Empire version.
        """
        return jsonify({'version': empire.VERSION})

    @app.route('/api/map', methods=['GET'])
    def list_routes():
        """
        List all of the current registered API routes.
        """
        output = {}
        for rule in app.url_map.iter_rules():
            methods = ','.join(rule.methods)
            url = rule.rule
            output.update({rule.endpoint: {'methods': methods, 'url': url}})

        return jsonify({'Routes':output})

    @app.route('/api/config', methods=['GET'])
    def get_config():
        """
        Returns JSON of the current Empire config.
        """
        api_username = g.user['username']
        api_current_token = g.user['api_token']
        configRaw = execute_db_query(conn, 'SELECT staging_key, install_path, ip_whitelist, ip_blacklist, autorun_command, autorun_data, rootuser FROM config')

        [staging_key, install_path, ip_whitelist, ip_blacklist, autorun_command, autorun_data, rootuser] = configRaw[0]
        config = [{"api_username":api_username, "autorun_command":autorun_command, "autorun_data":autorun_data, "current_api_token":api_current_token, "install_path":install_path, "ip_blacklist":ip_blacklist, "ip_whitelist":ip_whitelist, "staging_key":staging_key, "version":empire.VERSION}]
        return jsonify({'config': config})

    @app.route('/api/stagers', methods=['GET'])
    def get_stagers():
        """
        Returns JSON describing all stagers.
        """

        stagers = []
        for stagerName, stager in main.stagers.stagers.items():
            info = copy.deepcopy(stager.info)
            info['options'] = stager.options
            info['Name'] = stagerName
            stagers.append(info)

        return jsonify({'stagers': stagers})

    @app.route('/api/stagers/<path:stager_name>', methods=['GET'])
    def get_stagers_name(stager_name):
        """
        Returns JSON describing the specified stager_name passed.
        """
        if stager_name not in main.stagers.stagers:
            return make_response(jsonify({'error': 'stager name %s not found, make sure to use [os]/[name] format, ie. windows/dll' %(stager_name)}), 404)

        stagers = []
        for stagerName, stager in main.stagers.stagers.items():
            if stagerName == stager_name:
                info = copy.deepcopy(stager.info)
                info['options'] = stager.options
                info['Name'] = stagerName
                stagers.append(info)

        return jsonify({'stagers': stagers})

    @app.route('/api/stagers', methods=['POST'])
    def generate_stager():
        """
        Generates a stager with the supplied config and returns JSON information
        describing the generated stager, with 'Output' being the stager output.

        Required JSON args:
            StagerName      -   the stager name to generate
            Listener        -   the Listener name to use for the stager
        """
        if not request.json or not 'StagerName' in request.json or not 'Listener' in request.json:
            abort(400)

        stagerName = request.json['StagerName']
        listener = request.json['Listener']

        if stagerName not in main.stagers.stagers:
            return make_response(jsonify({'error': 'stager name %s not found' %(stagerName)}), 404)

        if not main.listeners.is_listener_valid(listener):
            return make_response(jsonify({'error': 'invalid listener ID or name'}), 400)

        stager = main.stagers.stagers[stagerName]

        # set all passed options
        for option, values in request.json.items():
            if option != 'StagerName':
                if option not in stager.options:
                    return make_response(jsonify({'error': 'Invalid option %s, check capitalization.' %(option)}), 400)
                stager.options[option]['Value'] = values

        # validate stager options
        for option, values in stager.options.items():
            if values['Required'] and ((not values['Value']) or (values['Value'] == '')):
                return make_response(jsonify({'error': 'required stager options missing'}), 400)

        stagerOut = copy.deepcopy(stager.options)

        if ('OutFile' in stagerOut) and (stagerOut['OutFile']['Value'] != ''):
            if isinstance(stager.generate(), str):
                # if the output was intended for a file, return the base64 encoded text
                stagerOut['Output'] = base64.b64encode(stager.generate().encode('UTF-8'))
            else:
                stagerOut['Output'] = base64.b64encode(stager.generate())

        else:
            # otherwise return the text of the stager generation
            stagerOut['Output'] = stager.generate()

        return jsonify({stagerName: stagerOut})

    @app.route('/api/modules', methods=['GET'])
    def get_modules():
        """
        Returns JSON describing all currently loaded modules.
        """

        modules = []
        for moduleName, module in main.modules.modules.items():
            moduleInfo = copy.deepcopy(module.info)
            moduleInfo['options'] = module.options
            moduleInfo['Name'] = moduleName
            modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/<path:module_name>', methods=['GET'])
    def get_module_name(module_name):
        """
        Returns JSON describing the specified currently module.
        """

        if module_name not in main.modules.modules:
            return make_response(jsonify({'error': 'module name %s not found' %(module_name)}), 404)

        modules = []
        moduleInfo = copy.deepcopy(main.modules.modules[module_name].info)
        moduleInfo['options'] = main.modules.modules[module_name].options
        moduleInfo['Name'] = module_name
        modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/<path:module_name>', methods=['POST'])
    def execute_module(module_name):
        """
        Executes a given module name with the specified parameters.
        """

        # ensure the 'Agent' argument is set
        if not request.json or not 'Agent' in request.json:
            abort(400)

        if module_name not in main.modules.modules:
            return make_response(jsonify({'error': 'module name %s not found' %(module_name)}), 404)

        module = main.modules.modules[module_name]

        # set all passed module options
        for key, value in request.json.items():
            if key not in module.options:
                return make_response(jsonify({'error': 'invalid module option'}), 400)

            module.options[key]['Value'] = value

        # validate module options
        sessionID = module.options['Agent']['Value']

        for option, values in module.options.items():
            if values['Required'] and ((not values['Value']) or (values['Value'] == '')):
                return make_response(jsonify({'error': 'required module option missing'}), 400)

        try:
            # if we're running this module for all agents, skip this validation
            if sessionID.lower() != "all" and sessionID.lower() != "autorun":

                if not main.agents.is_agent_present(sessionID):
                    return make_response(jsonify({'error': 'invalid agent name'}), 400)

                moduleVersion = float(module.info['MinLanguageVersion'])
                agentVersion = float(main.agents.get_language_version_db(sessionID))
                # check if the agent/module PowerShell versions are compatible
                if moduleVersion > agentVersion:
                    return make_response(jsonify({'error': "module requires PS version "+str(modulePSVersion)+" but agent running PS version "+str(agentPSVersion)}), 400)

        except Exception as e:
            return make_response(jsonify({'error': 'exception: %s' %(e)}), 400)

        # check if the module needs admin privs
        if module.info['NeedsAdmin']:
            # if we're running this module for all agents, skip this validation
            if sessionID.lower() != "all" and sessionID.lower() != "autorun":
                if not main.agents.is_agent_elevated(sessionID):
                    return make_response(jsonify({'error': 'module needs to run in an elevated context'}), 400)

        # actually execute the module
        moduleData = module.generate()

        if not moduleData or moduleData == "":
            return make_response(jsonify({'error': 'module produced an empty script'}), 400)

        try:
            if isinstance(moduleData, bytes):
                moduleData = moduleData.decode('ascii')
        except UnicodeDecodeError:
            return make_response(jsonify({'error': 'module source contains non-ascii characters'}), 400)

        moduleData = helpers.strip_powershell_comments(moduleData)
        taskCommand = ""

        # build the appropriate task command and module data blob
        if str(module.info['Background']).lower() == "true":
            # if this module should be run in the background
            extention = module.info['OutputExtension']
            if extention and extention != "":
                # if this module needs to save its file output to the server
                #   format- [15 chars of prefix][5 chars extension][data]
                saveFilePrefix = module_name.split("/")[-1]
                moduleData = saveFilePrefix.rjust(15) + extention.rjust(5) + moduleData
                taskCommand = "TASK_CMD_JOB_SAVE"
            else:
                taskCommand = "TASK_CMD_JOB"

        else:
            # if this module is run in the foreground
            extention = module.info['OutputExtension']
            if module.info['OutputExtension'] and module.info['OutputExtension'] != "":
                # if this module needs to save its file output to the server
                #   format- [15 chars of prefix][5 chars extension][data]
                saveFilePrefix = module_name.split("/")[-1][:15]
                moduleData = saveFilePrefix.rjust(15) + extention.rjust(5) + moduleData
                taskCommand = "TASK_CMD_WAIT_SAVE"
            else:
                taskCommand = "TASK_CMD_WAIT"

        if sessionID.lower() == "all":

            for agent in main.agents.get_agents():
                sessionID = agent[1]
                taskID = main.agents.add_agent_task_db(sessionID, taskCommand, moduleData, moduleName=module_name, uid=g.user['id'])
                msg = "tasked agent %s to run module %s" % (sessionID, module_name)
                main.agents.save_agent_log(sessionID, msg)

            msg = "tasked all agents to run module %s" %(module_name)
            return jsonify({'success': True, 'taskID': taskID, 'msg':msg})

        else:
            # set the agent's tasking in the cache
            taskID = main.agents.add_agent_task_db(sessionID, taskCommand, moduleData, moduleName=module_name, uid=g.user['id'])

            # update the agent log
            msg = "tasked agent %s to run module %s" %(sessionID, module_name)
            main.agents.save_agent_log(sessionID, msg)
            return jsonify({'success': True, 'taskID': taskID, 'msg':msg})

    @app.route('/api/modules/search', methods=['POST'])
    def search_modules():
        """
        Returns JSON describing the the modules matching the passed
        'term' search parameter. Module name, description, comments,
        and author fields are searched.
        """

        if not request.json or not 'term':
            abort(400)

        searchTerm = request.json['term']

        modules = []

        for moduleName, module in main.modules.modules.items():
            if (searchTerm.lower() == '') or (searchTerm.lower() in moduleName.lower()) or (
                    searchTerm.lower() in ("".join(module.info['Description'])).lower()) or (
                    searchTerm.lower() in ("".join(module.info['Comments'])).lower()) or (
                    searchTerm.lower() in ("".join(module.info['Author'])).lower()):
                moduleInfo = copy.deepcopy(main.modules.modules[moduleName].info)
                moduleInfo['options'] = main.modules.modules[moduleName].options
                moduleInfo['Name'] = moduleName
                modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/search/modulename', methods=['POST'])
    def search_modules_name():
        """
        Returns JSON describing the the modules matching the passed
        'term' search parameter for the modfule name.
        """

        if not request.json or not 'term':
            abort(400)

        searchTerm = request.json['term']

        modules = []

        for moduleName, module in main.modules.modules.items():
            if (searchTerm.lower() == '') or (searchTerm.lower() in moduleName.lower()):
                moduleInfo = copy.deepcopy(main.modules.modules[moduleName].info)
                moduleInfo['options'] = main.modules.modules[moduleName].options
                moduleInfo['Name'] = moduleName
                modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/search/description', methods=['POST'])
    def search_modules_description():
        """
        Returns JSON describing the the modules matching the passed
        'term' search parameter for the 'Description' field.
        """

        if not request.json or not 'term':
            abort(400)

        searchTerm = request.json['term']

        modules = []

        for moduleName, module in main.modules.modules.items():
            if (searchTerm.lower() == '') or (searchTerm.lower() in ("".join(module.info['Description'])).lower()):
                moduleInfo = copy.deepcopy(main.modules.modules[moduleName].info)
                moduleInfo['options'] = main.modules.modules[moduleName].options
                moduleInfo['Name'] = moduleName
                modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/search/comments', methods=['POST'])
    def search_modules_comments():
        """
        Returns JSON describing the the modules matching the passed
        'term' search parameter for the 'Comments' field.
        """

        if not request.json or not 'term':
            abort(400)

        searchTerm = request.json['term']

        modules = []

        for moduleName, module in main.modules.modules.items():
            if (searchTerm.lower() == '') or (searchTerm.lower() in ("".join(module.info['Comments'])).lower()):
                moduleInfo = copy.deepcopy(main.modules.modules[moduleName].info)
                moduleInfo['options'] = main.modules.modules[moduleName].options
                moduleInfo['Name'] = moduleName
                modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/modules/search/author', methods=['POST'])
    def search_modules_author():
        """
        Returns JSON describing the the modules matching the passed
        'term' search parameter for the 'Author' field.
        """

        if not request.json or not 'term':
            abort(400)

        searchTerm = request.json['term']

        modules = []

        for moduleName, module in main.modules.modules.items():
            if (searchTerm.lower() == '') or (searchTerm.lower() in ("".join(module.info['Author'])).lower()):
                moduleInfo = copy.deepcopy(main.modules.modules[moduleName].info)
                moduleInfo['options'] = main.modules.modules[moduleName].options
                moduleInfo['Name'] = moduleName
                modules.append(moduleInfo)

        return jsonify({'modules': modules})

    @app.route('/api/listeners', methods=['GET'])
    def get_listeners():
        """
        Returns JSON describing all currently registered listeners.
        """
        activeListenersRaw = execute_db_query(conn,
                                              'SELECT id, name, module, listener_type, listener_category, options, created_at FROM listeners')
        listeners = []

        for activeListener in activeListenersRaw:
            [ID, name, module, listener_type, listener_category, options, created_at] = activeListener
            listeners.append({'ID': ID, 'name': name, 'module': module, 'listener_type': listener_type,
                              'listener_category': listener_category, 'options': pickle.loads(options),
                              'created_at': created_at})

        return jsonify({'listeners': listeners})

    @app.route('/api/listeners/<string:listener_name>', methods=['GET'])
    def get_listener_name(listener_name):
        """
        Returns JSON describing the listener specified by listener_name.
        """
        activeListenersRaw = execute_db_query(conn,
                                              'SELECT id, name, module, listener_type, listener_category, options FROM listeners WHERE name=?',
                                              [listener_name])
        listeners = []

        # if listener_name != "" and main.listeners.is_listener_valid(listener_name):
        for activeListener in activeListenersRaw:
            [ID, name, module, listener_type, listener_category, options] = activeListener
            if name == listener_name:
                listeners.append({'ID': ID, 'name': name, 'module': module, 'listener_type': listener_type,
                                  'listener_category': listener_category, 'options': pickle.loads(activeListener[5])})

            return jsonify({'listeners': listeners})
        else:
            return make_response(jsonify({'error': 'listener name %s not found' % (listener_name)}), 404)

    @app.route('/api/listeners/<string:listener_name>', methods=['DELETE'])
    def kill_listener(listener_name):
        """
        Kills the listener specified by listener_name.
        """
        if listener_name.lower() == "all":
            activeListenersRaw = execute_db_query(conn, 'SELECT id, name, module, listener_type, listener_category, options FROM listeners')
            for activeListener in activeListenersRaw:
                [ID, name, module, listener_type, listener_category, options] = activeListener
                main.listeners.kill_listener(name)

            return jsonify({'success': True})
        else:
            if listener_name != "" and main.listeners.is_listener_valid(listener_name):
                main.listeners.kill_listener(listener_name)
                return jsonify({'success': True})
            else:
                return make_response(jsonify({'error': 'listener name %s not found' %(listener_name)}), 404)

    @app.route('/api/listeners/types', methods=['GET'])
    def get_listener_types():
        """
        Returns a list of the loaded listeners that are available for use.
        """

        return jsonify({'types' : list(main.listeners.loadedListeners.keys())})

    @app.route('/api/listeners/options/<string:listener_type>', methods=['GET'])
    def get_listener_options(listener_type):
        """
        Returns JSON describing listener options for the specified listener type.
        """

        if listener_type.lower() not in main.listeners.loadedListeners:
            return make_response(jsonify({'error':'listener type %s not found' %(listener_type)}), 404)

        options = main.listeners.loadedListeners[listener_type].options
        return jsonify({'listeneroptions' : options})

    @app.route('/api/listeners/<string:listener_type>', methods=['POST'])
    def start_listener(listener_type):
        """
        Starts a listener with options supplied in the POST.
        """
        if listener_type.lower() not in main.listeners.loadedListeners:
            return make_response(jsonify({'error':'listener type %s not found' %(listener_type)}), 404)

        listenerObject = main.listeners.loadedListeners[listener_type]
        # set all passed options
        for option, values in request.json.items():
            if isinstance(values, bytes):
                values = values.decode('UTF-8')
            if option == "Name":
                listenerName = values

            returnVal = main.listeners.set_listener_option(listener_type, option, values)
            if not returnVal:
                 return make_response(jsonify({'error': 'error setting listener value %s with option %s' %(option, values)}), 400)

        main.listeners.start_listener(listener_type, listenerObject)

        # check to see if the listener was created
        listenerID = main.listeners.get_listener_id(listenerName)
        if listenerID:
            return jsonify({'success': 'Listener %s successfully started' % listenerName})
        else:
            return jsonify({'error': 'failed to start listener %s' % listenerName})

    @app.route('/api/agents', methods=['GET'])
    def get_agents():
        """
        Returns JSON describing all currently registered agents.
        """
        activeAgentsRaw = execute_db_query(conn, 'SELECT id, session_id, listener, name, language, language_version, delay, jitter, external_ip, '+
            'internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, '+
            'lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results, notes FROM agents')
        agents = []

        for activeAgent in activeAgentsRaw:
            [ID, session_id, listener, name, language, language_version, delay, jitter, external_ip, internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results, notes] = activeAgent

            stale = helpers.is_stale(lastseen_time, delay, jitter)

            if isinstance(session_key, bytes):
                session_key = session_key.decode('latin-1').encode('utf-8')

            agents.append({"ID":ID, "session_id":session_id, "listener":listener, "name":name, "language":language, "language_version":language_version, "delay":delay, "jitter":jitter, "external_ip":external_ip, "internal_ip":internal_ip, "username":username, "high_integrity":high_integrity, "process_name":process_name, "process_id":process_id, "hostname":hostname, "os_details":os_details, "session_key":session_key, "nonce":nonce, "checkin_time":checkin_time, "lastseen_time":lastseen_time, "parent":parent, "children":children, "servers":servers, "profile":profile,"functions":functions, "kill_date":kill_date, "working_hours":working_hours, "lost_limit":lost_limit, "taskings":taskings, "results":results, "stale":stale, "notes":notes})

        return jsonify({'agents' : agents})

    @app.route('/api/agents/stale', methods=['GET'])
    def get_agents_stale():
        """
        Returns JSON describing all stale agents.
        """

        agentsRaw = execute_db_query(conn, 'SELECT id, session_id, listener, name, language, language_version, delay, jitter, external_ip, '+
            'internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, '+
            'lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results FROM agents')
        staleAgents = []

        for agent in agentsRaw:
            [ID, session_id, listener, name, language, language_version, delay, jitter, external_ip, internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results] = agent

            stale = helpers.is_stale(lastseen_time, delay, jitter)

            if stale:
                if isinstance(session_key, bytes):
                    session_key = session_key.decode('latin-1').encode('utf-8')

                staleAgents.append({"ID":ID, "session_id":session_id, "listener":listener, "name":name, "language":language, "language_version":language_version, "delay":delay, "jitter":jitter, "external_ip":external_ip, "internal_ip":internal_ip, "username":username, "high_integrity":high_integrity, "process_name":process_name, "process_id":process_id, "hostname":hostname, "os_details":os_details, "session_key":session_key, "nonce":nonce, "checkin_time":checkin_time, "lastseen_time":lastseen_time, "parent":parent, "children":children, "servers":servers, "profile":profile,"functions":functions, "kill_date":kill_date, "working_hours":working_hours, "lost_limit":lost_limit, "taskings":taskings, "results":results})

        return jsonify({'agents' : staleAgents})

    @app.route('/api/agents/stale', methods=['DELETE'])
    def remove_stale_agent():
        """
        Removes stale agents from the controller.

        WARNING: doesn't kill the agent first! Ensure the agent is dead.
        """
        agentsRaw = execute_db_query(conn, 'SELECT * FROM agents')

        for agent in agentsRaw:
            [ID, sessionID, listener, name, language, language_version, delay, jitter, external_ip, internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results] = agent

            stale = helpers.is_stale(lastseen_time, delay, jitter)

            if stale:
                execute_db_query(conn, "DELETE FROM agents WHERE session_id LIKE ?", [sessionID])

        return jsonify({'success': True})

    @app.route('/api/agents/<string:agent_name>', methods=['DELETE'])
    def remove_agent(agent_name):
        """
        Removes an agent from the controller specified by agent_name.

        WARNING: doesn't kill the agent first! Ensure the agent is dead.
        """
        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            execute_db_query(conn, "DELETE FROM agents WHERE session_id LIKE ?", [agentSessionID])

        return jsonify({'success': True})

    @app.route('/api/agents/<string:agent_name>', methods=['GET'])
    def get_agents_name(agent_name):
        """
        Returns JSON describing the agent specified by agent_name.
        """
        activeAgentsRaw = execute_db_query(conn, 'SELECT id, session_id, listener, name, language, language_version, delay, jitter, external_ip, '+
            'internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, '+
            'lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results FROM agents ' +
            'WHERE name=? OR session_id=?', [agent_name, agent_name])
        activeAgents = []

        for activeAgent in activeAgentsRaw:
            [ID, session_id, listener, name, language, language_version, delay, jitter, external_ip, internal_ip, username, high_integrity, process_name, process_id, hostname, os_details, session_key, nonce, checkin_time, lastseen_time, parent, children, servers, profile, functions, kill_date, working_hours, lost_limit, taskings, results] = activeAgent

            if isinstance(session_key, bytes):
                session_key = session_key.decode('latin-1').encode('utf-8')

            activeAgents.append({"ID":ID, "session_id":session_id, "listener":listener, "name":name, "language":language, "language_version":language_version, "delay":delay, "jitter":jitter, "external_ip":external_ip, "internal_ip":internal_ip, "username":username, "high_integrity":high_integrity, "process_name":process_name, "process_id":process_id, "hostname":hostname, "os_details":os_details, "session_key":session_key, "nonce":nonce, "checkin_time":checkin_time, "lastseen_time":lastseen_time, "parent":parent, "children":children, "servers":servers, "profile":profile,"functions":functions, "kill_date":kill_date, "working_hours":working_hours, "lost_limit":lost_limit, "taskings":taskings, "results":results})

        return jsonify({'agents' : activeAgents})

    @app.route('/api/agents/<string:agent_name>/directory', methods=['POST'])
    def scrape_agent_directory(agent_name):
        directory = '/' if request.args.get('directory') is None else request.args.get('directory')
        task_id = main.agents.add_agent_task_db(agent_name, "TASK_DIR_LIST", directory, g.user['id'])
        return jsonify({'taskID': task_id})

    @app.route('/api/agents/<string:agent_name>/directory', methods=['GET'])
    def get_agent_directory(agent_name):
        # Would be cool to add a "depth" param
        directory = '/' if request.args.get('directory') is None else request.args.get('directory')

        found = execute_db_query(conn, "SELECT * FROM file_directory WHERE session_id = ? AND path = ? AND is_file = 0",
                                 [agent_name, directory])
        if len(found) == 0:
            return make_response(jsonify({'error': "Directory not found."}), 404)

        results = execute_db_query(conn, """
            SELECT
                base.id,
                base.session_id,
                base.name,
                base.path,
                base.parent_id,
                base.is_file,
                p.name,
                p.path,
                p.parent_id
            FROM file_directory base
            join file_directory p on base.parent_id = p.id
            WHERE base.session_id = ? and p.path = ?
        """, [agent_name, directory])

        response = []
        for result in results:
            [id, session_id, name, path, parent_id, is_file, parent_name, parent_path, parent_parent] = result
            response.append({'id': id, 'session_id': session_id, 'name': name, 'path': path, 'parent_id': parent_id, 'is_file': bool(is_file),
                           'parent_name': parent_name, 'parent_path': parent_path, 'parent_parent': parent_parent})

        return jsonify({'items': response})

    @app.route('/api/agents/<string:agent_name>/results', methods=['GET'])
    def get_agent_results(agent_name):
        """
        Returns JSON describing the agent's results and removes the result field
        from the backend database.
        """
        agentTaskResults = []

        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name, session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name, session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        for agentNameID in agentNameIDs:
            [agentName, agentSessionID] = agentNameID

            agentResults = execute_db_query(conn,
                """
                SELECT
                    results.id,
                    taskings.data AS command,
                    results.data AS response,
                    users.id as user_id,
                    users.username as username
                FROM results
                INNER JOIN taskings ON results.id = taskings.id AND results.agent = taskings.agent
                LEFT JOIN users on results.user_id = users.id
                WHERE results.agent=?;
                """, [agentSessionID])

            results = []
            if len(agentResults) > 0:
                for taskID, command, result, user_id, username in agentResults:
                    results.append({'taskID': taskID, 'command': command, 'results': result, 'user_id': user_id, 'username': username})

                agentTaskResults.append({"AgentName": agentSessionID, "AgentResults": results})
            else:
                agentTaskResults.append({"AgentName": agentSessionID, "AgentResults": []})

        return jsonify({'results': agentTaskResults})

    @app.route('/api/agents/<string:agent_name>/task/<int:task_id>', methods=['GET'])
    def get_task(agent_name, task_id):
        results = execute_db_query(conn, """
            SELECT
                taskings.id AS task,
                taskings.data AS command,
                results.data AS response,
                users.id AS user_id,
                users.username AS username
             FROM taskings
             LEFT JOIN users ON taskings.user_id = users.id
             LEFT JOIN results on results.id = taskings.id AND results.agent = taskings.agent
             WHERE taskings.agent = ?
             AND taskings.id = ?
             """, [agent_name, task_id])

        if len(results) > 0:
            [taskID, command, results, user_id, username] = results[0]
            return make_response(jsonify({'taskID': taskID, 'command': command, 'results': results, 'user_id': user_id, 'username': username}))

        return make_response(jsonify({'error': 'task not found.'}), 404)

    @app.route('/api/agents/<string:agent_name>/results', methods=['DELETE'])
    def delete_agent_results(agent_name):
        """
        Removes the specified agent results field from the backend database.
        """
        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            execute_db_query(conn, 'UPDATE agents SET results=? WHERE session_id=? OR name=?', ['', agentSessionID, agentName])

        return jsonify({'success': True})

    @app.route('/api/agents/<string:agent_name>/download', methods=['POST'])
    def task_agent_download(agent_name):
        """
        Tasks the specified agent to download a file
        """

        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agent_name_ids = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agent_name_ids = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agent_name_ids or len(agent_name_ids) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' % agent_name}), 404)

        if not request.json['filename']:
            return make_response(jsonify({'error':'file name not provided'}), 404)

        file_name = request.json['filename']
        for agent_name_id in agent_name_ids:
            (agentName, agentSessionID) = agent_name_id

            msg = "Tasked agent to download %s" % file_name
            main.agents.save_agent_log(agentSessionID, msg)
            task_id = main.agents.add_agent_task_db(agentSessionID, 'TASK_DOWNLOAD', file_name, uid=g.user['id'])

        return jsonify({'success': True, 'taskID': task_id})

    @app.route('/api/agents/<string:agent_name>/upload', methods=['POST'])
    def task_agent_upload(agent_name):
        """
        Tasks the specified agent to upload a file
        """

        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        if not request.json['data']:
            return make_response(jsonify({'error':'file data not provided'}), 404)

        if not request.json['filename']:
            return make_response(jsonify({'error':'file name not provided'}), 404)

        fileData = request.json['data']
        fileName = request.json['filename']

        rawBytes = base64.b64decode(fileData)

        if len(rawBytes) > 1048576:
            return make_response(jsonify({'error':'file size too large'}), 404)

        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            msg = "Tasked agent to upload %s : %s" % (fileName, hashlib.md5(rawBytes).hexdigest())
            main.agents.save_agent_log(agentSessionID, msg)
            data = fileName + "|" + fileData
            taskID = main.agents.add_agent_task_db(agentSessionID, 'TASK_UPLOAD', data, uid=g.user['id'])

        return jsonify({'success': True, 'taskID': taskID})

    @app.route('/api/agents/<string:agent_name>/shell', methods=['POST'])
    def task_agent_shell(agent_name):
        """
        Tasks an the specified agent_name to execute a shell command.

        Takes {'command':'shell_command'}
        """
        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        command = request.json['command']
        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            # add task command to agent taskings
            msg = "tasked agent %s to run command %s" %(agentSessionID, command)
            main.agents.save_agent_log(agentSessionID, msg)
            taskID = main.agents.add_agent_task_db(agentSessionID, "TASK_SHELL", command, uid=g.user['id'])
        return jsonify({'success': True, 'taskID': taskID})

    @app.route('/api/agents/<string:agent_name>/update_comms', methods=['PUT'])
    def agent_update_comms(agent_name):
        """
        Dynamically update the agent comms to another

        Takes {'listener': 'name'}
        """

        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        if not 'listener' in request.json:
            return make_response(jsonify({'error':'JSON body must include key "listener"'}), 400)

        listener_name = request.json['listener']

        if not main.listeners.is_listener_valid(listener_name):
            return jsonify({'error': 'Please enter a valid listener name.'})
        else:
            active_listener = main.listeners.activeListeners[listener_name]
            if active_listener['moduleName'] != 'meterpreter' or active_listener['moduleName'] != 'http_mapi':
                listener_options = active_listener['options']
                listener_comms = main.listeners.loadedListeners[active_listener['moduleName']].generate_comms(listener_options, language="powershell")

                main.agents.add_agent_task_db(agent_name, "TASK_UPDATE_LISTENERNAME", listener_options['Name']['Value'])
                main.agents.add_agent_task_db(agent_name, "TASK_SWITCH_LISTENER", listener_comms)

                msg = "Tasked agent to update comms to %s listener" % listener_name
                main.agents.save_agent_log(agent_name, msg)
                return jsonify({'success': True})
            else:
                return jsonify({'error': 'Ineligible listener for updatecomms command: %s' % active_listener['moduleName']})

    @app.route('/api/agents/<string:agent_name>/killdate', methods=['PUT'])
    def agent_kill_date(agent_name):
        """
        Set an agent's killdate (01/01/2016)

        Takes {'kill_date': 'date'}
        """

        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        if not 'kill_date' in request.json:
            return make_response(jsonify({'error':'JSON body must include key "kill_date"'}), 400)

        try:
            kill_date = request.json['kill_date']

            # update this agent's information in the database
            main.agents.set_agent_field_db("kill_date", kill_date, agent_name)

            # task the agent
            main.agents.add_agent_task_db(agent_name, "TASK_SHELL", "Set-KillDate " + str(kill_date))

            # update the agent log
            msg = "Tasked agent to set killdate to " + str(kill_date)
            main.agents.save_agent_log(agent_name, msg)
            return jsonify({'success': True})
        except:
            return jsonify({'error': 'Unable to update agent killdate'})

    @app.route('/api/agents/<string:agent_name>/workinghours', methods=['PUT'])
    def agent_working_hours(agent_name):
        """
        Set an agent's working hours (9:00-17:00)

        Takes {'working_hours': 'working_hours'}
        """

        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        if not 'working_hours' in request.json:
            return make_response(jsonify({'error':'JSON body must include key "working_hours"'}), 400)

        try:
            working_hours = request.json['working_hours']
            working_hours = working_hours.replace(",", "-")

            # update this agent's information in the database
            main.agents.set_agent_field_db("working_hours", working_hours, agent_name)

            # task the agent
            main.agents.add_agent_task_db(agent_name, "TASK_SHELL", "Set-WorkingHours " + str(working_hours))

            # update the agent log
            msg = "Tasked agent to set working hours to " + str(working_hours)
            main.agents.save_agent_log(agent_name, msg)
            return jsonify({'success': True})
        except:
            return jsonify({'error': 'Unable to update agent workinghours'})

    @app.route('/api/agents/<string:agent_name>/rename', methods=['POST'])
    def task_agent_rename(agent_name):
        """
        Renames the specified agent.

        Takes {'newname': 'NAME'}
        """

        agentNameID = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameID or len(agentNameID) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        (agentName, agentSessionID) = agentNameID[0]
        newName = request.json['newname']

        try:
            result = main.agents.rename_agent(agentName, newName)

            if not result:
                return make_response(jsonify({'error': 'error in renaming %s to %s, new name may have already been used' %(agentName, newName)}), 400)

            return jsonify({'success': True})

        except Exception:
            return make_response(jsonify({'error': 'error in renaming %s to %s' %(agentName, newName)}), 400)

    @app.route('/api/agents/<string:agent_name>/clear', methods=['POST', 'GET'])
    def task_agent_clear(agent_name):
        """
        Clears the tasking buffer for the specified agent.
        """
        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            main.agents.clear_agent_tasks_db(agentSessionID)

        return jsonify({'success': True})

    @app.route('/api/agents/<string:agent_name>/kill', methods=['POST', 'GET'])
    def task_agent_kill(agent_name):
        """
        Tasks the specified agent to exit.
        """
        if agent_name.lower() == "all":
            # enumerate all target agent sessionIDs
            agentNameIDs = execute_db_query(conn, "SELECT name,session_id FROM agents WHERE name like '%' OR session_id like '%'")
        else:
            agentNameIDs = execute_db_query(conn, 'SELECT name,session_id FROM agents WHERE name like ? OR session_id like ?', [agent_name, agent_name])

        if not agentNameIDs or len(agentNameIDs) == 0:
            return make_response(jsonify({'error': 'agent name %s not found' %(agent_name)}), 404)

        for agentNameID in agentNameIDs:
            (agentName, agentSessionID) = agentNameID

            # task the agent to exit
            msg = "tasked agent %s to exit" %(agentSessionID)
            main.agents.save_agent_log(agentSessionID, msg)
            main.agents.add_agent_task_db(agentSessionID, 'TASK_EXIT', uid=g.user['id'])

        return jsonify({'success': True})

    @app.route('/api/agents/<string:agent_name>/notes', methods=['POST'])
    def update_agent_notes(agent_name):
        """
        Update notes on specified agent.
        {"notes" : "notes here"}
        """

        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        if not 'notes' in request.json:
            return make_response(jsonify({'error':'JSON body must include key "notes"'}), 400)

        notes = request.json['notes']

        try:
            cur = conn.cursor()
            cur.execute("UPDATE agents SET notes = ? WHERE name = ?", (notes, agent_name))

        finally:
            cur.close()

        return jsonify({'success': True})

    @app.route('/api/creds', methods=['GET'])
    def get_creds():
        """
        Returns JSON describing the credentials stored in the backend database.
        """
        credsRaw = execute_db_query(conn, 'SELECT ID, credtype, domain, username, password, host, os, sid, notes FROM credentials')
        creds = []

        for credRaw in credsRaw:
            [ID, credtype, domain, username, password, host, os, sid, notes] = credRaw
            creds.append({"ID":ID, "credtype":credtype, "domain":domain, "username":username, "password":password, "host":host, "os":os, "sid":sid, "notes":notes})

        return jsonify({'creds' : creds})

    @app.route('/api/creds', methods=['POST'])
    def add_creds():
        """
        Adds credentials to the database
        """
        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        if not 'credentials' in request.json:
            return make_response(jsonify({'error':'JSON body must include key "credentials"'}), 400)

        creds = request.json['credentials']

        if not type(creds) == list:
            return make_response(jsonify({'error':'credentials must be provided as a list'}), 400)

        required_fields = ["credtype", "domain", "username", "password", "host"]
        optional_fields = ["OS", "notes", "sid"]

        for cred in creds:
            # ensure every credential given to us has all the required fields
            if not all (k in cred for k in required_fields):
                return make_response(jsonify({'error':'invalid credential %s' %(cred)}), 400)

            # ensure the type is either "hash" or "plaintext"
            if not (cred['credtype'] == u'hash' or cred['credtype'] == u'plaintext'):
                return make_response(jsonify({'error':'invalid credential type in %s, must be "hash" or "plaintext"' %(cred)}), 400)

        # other than that... just assume everything is valid

        # this would be way faster if batched but will work for now
        for cred in creds:
            # get the optional stuff, if it's there
            try:
                os = cred['os']
            except KeyError:
                os = ''

            try:
                sid = cred['sid']
            except KeyError:
                sid = ''

            try:
                notes = cred['notes']
            except KeyError:
                notes = ''

            main.credentials.add_credential(
                cred['credtype'],
                cred['domain'],
                cred['username'],
                cred['password'],
                cred['host'],
                os,
                sid,
                notes
            )

        return jsonify({'success': '%s credentials added' % len(creds)})

    @app.route('/api/reporting', methods=['GET'])
    def get_reporting():
        """
        Returns JSON describing the reporting events from the backend database.
        """
        reportingRaw = execute_db_query(conn, '''
        SELECT
                reporting.timestamp,
                event_type,
                u.username,
                substr(reporting.name, pos+1) as agent_name,
                a.hostname,
                taskID,
                t.data as "Task",
                r.data as "Results"
            FROM
            (
                SELECT
                    timestamp,
                    event_type,
                    name,
                    instr(name, '/') as pos,
                    taskID
                FROM reporting
                WHERE name LIKE 'agent%'
                AND reporting.event_type == 'task' OR reporting.event_type == 'checkin'
            ) reporting
            LEFT OUTER JOIN taskings t on (reporting.taskID = t.id) AND (agent_name = t.agent)
            LEFT OUTER JOIN results r on (reporting.taskID = r.id) AND (agent_name = r.agent)
            JOIN agents a on agent_name = a.session_id
            LEFT OUTER JOIN users u on t.user_id = u.id
            ORDER BY reporting.timestamp DESC
        ''')
        reportingEvents = []

        for reportingEvent in reportingRaw:
            [timestamp, event_type, user_name, agent_name, host_name, taskID, task, results] = reportingEvent
            reportingEvents.append({"timestamp":timestamp, "event_type":event_type, "username":user_name, "agent_name":agent_name, "host_name":host_name, "taskID":taskID, "task":task, "results":results})

        return jsonify({'reporting' : reportingEvents})

    @app.route('/api/reporting/generate', methods=['POST'])
    def generate_report():
        """
        Generates reports on the backend database.
        Takes {'logo':'directory_location'}
        """
        if not request.json:
            return make_response(jsonify({'error':'request body must be valid JSON'}), 400)

        directory_location = request.json['logo']

        main.do_report(directory_location)
        return jsonify({'success': True})

    @app.route('/api/reporting/agent/<string:reporting_agent>', methods=['GET'])
    def get_reporting_agent(reporting_agent):
        """
        Returns JSON describing the reporting events from the backend database for
        the agent specified by reporting_agent.
        """

        # first resolve the supplied name to a sessionID
        results = execute_db_query(conn, 'SELECT session_id FROM agents WHERE name=?', [reporting_agent])
        if results:
            sessionID = results[0][0]
        else:
            return jsonify({'reporting' : ''})

        reportingRaw = execute_db_query(conn, 'SELECT ID, name, event_type, message, timestamp, taskID FROM reporting WHERE name=?', [sessionID])
        reportingEvents = []

        for reportingEvent in reportingRaw:
            [ID, name, event_type, message, timestamp, taskID] = reportingEvent
            reportingEvents.append({"ID":ID, "agentname":name, "event_type":event_type, "message":json.loads(message), "timestamp":timestamp, "taskID":taskID})

        return jsonify({'reporting' : reportingEvents})

    @app.route('/api/reporting/type/<string:event_type>', methods=['GET'])
    def get_reporting_type(event_type):
        """
        Returns JSON describing the reporting events from the backend database for
        the event type specified by event_type.
        """
        reportingRaw = execute_db_query(conn, 'SELECT ID, name, event_type, message, timestamp, taskID FROM reporting WHERE event_type=?', [event_type])
        reportingEvents = []

        for reportingEvent in reportingRaw:
            [ID, name, event_type, message, timestamp, taskID] = reportingEvent
            reportingEvents.append({"ID":ID, "agentname":name, "event_type":event_type, "message":json.loads(message), "timestamp":timestamp, "taskID":taskID})

        return jsonify({'reporting' : reportingEvents})

    @app.route('/api/reporting/msg/<string:msg>', methods=['GET'])
    def get_reporting_msg(msg):
        """
        Returns JSON describing the reporting events from the backend database for
        the any messages with *msg* specified by msg.
        """
        reportingRaw = execute_db_query(conn, "SELECT ID, name, event_type, message, timestamp, taskID FROM reporting WHERE message like ?", ['%'+msg+'%'])
        reportingEvents = []

        for reportingEvent in reportingRaw:
            [ID, name, event_type, message, timestamp, taskID] = reportingEvent
            reportingEvents.append({"ID":ID, "agentname":name, "event_type":event_type, "message":json.loads(message), "timestamp":timestamp, "taskID":taskID})

        return jsonify({'reporting' : reportingEvents})

    @app.route('/api/admin/login', methods=['POST'])
    def server_login():
        """
        Takes a supplied username and password and returns the current API token
        if authentication is accepted.
        """
        if not request.json or not 'username' in request.json or not 'password' in request.json:
            abort(400)

        suppliedUsername = request.json['username']
        suppliedPassword = request.json['password']

        # try to prevent some basic bruting
        time.sleep(2)
        token = main.users.user_login(suppliedUsername, suppliedPassword)

        if token:
            return jsonify({'token': token})
        else:
            return make_response('', 401)

    @app.route('/api/admin/logout', methods=['POST'])
    def server_logout():
        """
        Logs out current user
        """
        main.users.user_logout(g.user['id'])
        return jsonify({'success': True})

    @app.route('/api/admin/restart', methods=['GET', 'POST', 'PUT'])
    def signal_server_restart():
        """
        Signal a restart for the Flask server and any Empire instance.
        """
        restart_server()
        return jsonify({'success': True})

    @app.route('/api/admin/shutdown', methods=['GET', 'POST', 'PUT'])
    def signal_server_shutdown():
        """
        Signal a restart for the Flask server and any Empire instance.
        """
        shutdown_server()
        return jsonify({'success': True})

    @app.route('/api/users', methods=['GET'])
    def get_users():
        """
        Returns JSON of the users from the backend database.
        """
        reportingRaw = execute_db_query(conn, 'SELECT ID, username, last_logon_time, enabled, admin FROM users')
        reporting_users = []
        user_report = []

        for reporting_users in reportingRaw:
            [ID, username, last_logon_time, enabled, admin] = reporting_users
            data = {"ID": ID, "username": username, "last_logon_time": last_logon_time, "enabled": bool(enabled), "admin": bool(admin)}
            user_report.append(data)

        return jsonify({'users': user_report})

    @app.route('/api/users/<int:uid>', methods=['GET'])
    def get_user(uid):
        """
        return the user for an id
        """
        user = execute_db_query(conn, 'SELECT ID, username, last_logon_time, enabled, admin, notes FROM users WHERE id = ?', [uid,])

        if len(user) == 0:
            make_response(jsonify({'error': 'user %s not found' % uid}), 404)

        [ID, username, last_logon_time, enabled, admin, notes] = user[0]
        return jsonify({"ID": ID, "username": username, "last_logon_time": last_logon_time, "enabled": bool(enabled), "admin": bool(admin), "notes": notes})


    @app.route('/api/users/me', methods=['GET'])
    def get_user_me():
        """
        Returns the current user.
        """
        return jsonify(g.user)

    @app.route('/api/users', methods=['POST'])
    def create_user():
        # Check that input is a valid request
        if not request.json or not 'username' in request.json or not 'password' in request.json:
            abort(400)

        # Check if user is an admin
        if not main.users.is_admin(g.user['id']):
            abort(403)

        status = main.users.add_new_user(request.json['username'], request.json['password'])
        return jsonify({'success': status})

    @app.route('/api/users/<int:uid>/disable', methods=['PUT'])
    def disable_user(uid):
        # Don't disable yourself
        if not request.json or not 'disable' in request.json or uid == g.user['id']:
            abort(400)

        # User performing the action should be an admin.
        # User being updated should not be an admin.
        if not main.users.is_admin(g.user['id']) or main.users.is_admin(uid):
            abort(403)

        status = main.users.disable_user(uid, request.json['disable'])
        return jsonify({'success': status})

    @app.route('/api/users/<int:uid>/updatepassword', methods=['PUT'])
    def update_user_password(uid):
        if not request.json or not 'password' in request.json:
            abort(400)

        # Must be an admin or updating self.
        if not (main.users.is_admin(g.user['id']) or uid == g.user['id']):
            abort(403)

        status = main.users.update_password(uid, request.json['password'])
        return jsonify({'success': status})

    @app.route('/api/users/<int:uid>/notes', methods=['POST'])
    def update_user_notes(uid):
        """
        Update notes for a user.
        {"notes" : "notes here"}
        """

        if not request.json:
            return make_response(jsonify({'error': 'request body must be valid JSON'}), 400)

        if not 'notes' in request.json:
            return make_response(jsonify({'error': 'JSON body must include key "credentials"'}), 400)

        notes = request.json['notes']

        try:
            cur = conn.cursor()
            cur.execute("UPDATE users SET notes = ? WHERE id = ?", (notes, uid))

        finally:
            cur.close()

        return jsonify({'success': True})

    @app.route('/api/plugin/active', methods=['GET'])
    def list_active_plugins():
        """
        Lists all active plugins
        """
        plugins = []

        plugin_path = os.path.abspath("plugins")
        all_plugin_names = [name for _, name, _ in pkgutil.walk_packages([plugin_path])]
        # check if the plugin has already been loaded
        active_plugins = list(empireMenu.loadedPlugins.keys())
        for plugin_name in all_plugin_names:
            if plugin_name in active_plugins:
                data = empireMenu.loadedPlugins[plugin_name].info[0]
                data['options'] = empireMenu.loadedPlugins[plugin_name].options
                plugins.append(data)

        return jsonify({'plugins': plugins})

    @app.route('/api/plugin/<string:plugin_name>', methods=['GET'])
    def get_plugin(plugin_name):
        # check if the plugin has already been loaded
        if plugin_name not in empireMenu.loadedPlugins.keys():
            try:
                empireMenu.do_plugin(plugin_name)
            except:
                return make_response(jsonify({'error': 'plugin %s not found' % (plugin_name)}), 400)
        # get the commands available to the user. This can probably be done in one step if desired
        name = empireMenu.loadedPlugins[plugin_name].get_commands()['name']
        commands = empireMenu.loadedPlugins[plugin_name].get_commands()['commands']
        description = empireMenu.loadedPlugins[plugin_name].get_commands()['description']
        data = {'name': name, 'commands': commands, 'description': description}
        return jsonify(data)

    @app.route('/api/plugin/<string:plugin_name>', methods=['POST'])
    def execute_plugin(plugin_name):
        # check if the plugin has been loaded
        if plugin_name not in empireMenu.loadedPlugins.keys():
            return make_response(jsonify({'error': 'plugin %s not loaded' % (plugin_name)}), 404)

        use_plugin = empireMenu.loadedPlugins[plugin_name]

        # set all passed module options
        for key, value in request.json.items():
            if key not in use_plugin.options:
                return make_response(jsonify({'error': 'invalid module option'}), 400)

            use_plugin.options[key]['Value'] = value

        for option, values in use_plugin.options.items():
            if values['Required'] and ((not values['Value']) or (values['Value'] == '')):
                return make_response(jsonify({'error': 'required module option missing'}), 400)

        results = use_plugin.execute(request.json)
        if results == False:
            return make_response(jsonify({'error': 'internal plugin error'}), 400)
        return jsonify(results)

    if not os.path.exists('./data/empire-chain.pem'):
        print("[!] Error: cannot find certificate ./data/empire-chain.pem")
        sys.exit()

    def shutdown_server():
        """
        Shut down the Flask server and any Empire instance gracefully.
        """
        global serverExitCommand

        if suppress:
            # repair stdout
            sys.stdout.close()
            sys.stdout = oldStdout

        print("\n * Shutting down Empire RESTful API")

        if conn:
            conn.close()

        if suppress:
            print(" * Shutting down the Empire instance")
            main.shutdown()

        serverExitCommand = 'shutdown'

        func = request.environ.get('werkzeug.server.shutdown')
        if func is not None:
            func()

    def restart_server():
        """
        Restart the Flask server and any Empire instance.
        """
        global serverExitCommand

        shutdown_server()

        serverExitCommand = 'restart'

    def signal_handler(signal, frame):
        """
        Overrides the keyboardinterrupt signal handler so we can gracefully shut everything down.
        """

        global serverExitCommand

        with app.test_request_context():
            shutdown_server()

        serverExitCommand = 'shutdown'

        # repair the original signal handler
        import signal
        signal.signal(signal.SIGINT, signal.default_int_handler)
        sys.exit()

    try:
        signal.signal(signal.SIGINT, signal_handler)
    except ValueError:
        pass

    # wrap the Flask connection in SSL and start it
    certPath = os.path.abspath("./data/")

    # support any version of tls
    pyversion = sys.version_info
    if pyversion[0] == 2 and pyversion[1] == 7 and pyversion[2] >= 13:
        proto = ssl.PROTOCOL_TLS
    elif pyversion[0] >= 3:
        proto = ssl.PROTOCOL_TLS
    else:
        proto = ssl.PROTOCOL_SSLv23

    context = ssl.SSLContext(proto)
    context.load_cert_chain("%s/empire-chain.pem" % (certPath), "%s/empire-priv.key" % (certPath))
    app.run(host='0.0.0.0', port=int(port), ssl_context=context, threaded=True)


def start_sockets(empire_menu: MainMenu, port: int = 5000):
    app = Flask(__name__)
    socketio = SocketIO(app, cors_allowed_origins="*")
    empire_menu.socketio = socketio

    @socketio.on('connect')
    def connect():
        token = request.args.get('token', '')
        if len(token) > 0:
            user = empire_menu.users.get_user_from_token(token)
            if user:
                print(f"{user['username']} connected to socketio")
                return

        raise ConnectionRefusedError('unauthorized!')

    @socketio.on('disconnect')
    def test_disconnect():
        print('Client disconnected from socketio')

    print('')
    print(" * Starting Empire SocketIO on port: {}".format(port))

    cert_path = os.path.abspath("./data/")
    proto = ssl.PROTOCOL_TLS
    context = ssl.SSLContext(proto)
    context.load_cert_chain("{}/empire-chain.pem".format(cert_path), "{}/empire-priv.key".format(cert_path))
    socketio.run(app, host='0.0.0.0', port=port, ssl_context=context)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    generalGroup = parser.add_argument_group('General Options')
    generalGroup.add_argument('--debug', nargs='?', const='1', help='Debug level for output (default of 1, 2 for msg display).')
    generalGroup.add_argument('--reset', action='store_true', help="Resets Empire's database to defaults.")
    generalGroup.add_argument('-v', '--version', action='store_true', help='Display current Empire version.')
    generalGroup.add_argument('-r','--resource', nargs=1, help='Run the Empire commands in the specified resource file after startup.')

    cliGroup = parser.add_argument_group('CLI Payload Options')
    cliGroup.add_argument('-l', '--listener', nargs='?', const="list", help='Display listener options. Displays all listeners if nothing is specified.')
    cliGroup.add_argument('-s', '--stager', nargs='?', const="list", help='Specify a stager to generate. Lists all stagers if none is specified.')
    cliGroup.add_argument('-o', '--stager-options', nargs='*', help="Supply options to set for a stager in OPTION=VALUE format. Lists options if nothing is specified.")

    restGroup = parser.add_argument_group('RESTful API Options')
    launchGroup = restGroup.add_mutually_exclusive_group()
    launchGroup.add_argument('--rest', action='store_true', help='Run Empire and the RESTful API.')
    launchGroup.add_argument('--headless', action='store_true', help='Run Empire and the RESTful API headless without the usual interface.')
    restGroup.add_argument('--restport', type=int, nargs=1, help='Port to run the Empire RESTful API on. Defaults to 1337')
    restGroup.add_argument('-n', '--notifications', action='store_true', help='Run the SocketIO notifications server.')
    restGroup.add_argument('--socketport', type=int, nargs=1, help='Port to run socketio on. Defaults to 5000')
    restGroup.add_argument('--username', nargs=1, help='Start the RESTful API with the specified username instead of pulling from empire.db')
    restGroup.add_argument('--password', nargs=1, help='Start the RESTful API with the specified password instead of pulling from empire.db')

    args = parser.parse_args()
    database_check_docker()

    if not args.restport:
        args.restport = '1337'
    else:
        args.restport = args.restport[0]

    if not args.socketport:
        args.socketport = '5000'
    else:
        args.socketport = args.socketport[0]

    if args.version:
        print(empire.VERSION)

    if args.reset:
        choice = input("\n [>] Would you like to reset your Empire instance? [y/N]: ")
        if choice.lower() == "y":
            subprocess.call("./setup/reset.sh")
        else:
            pass

    elif args.rest:
        # start an Empire instance and RESTful API
        main = empire.MainMenu(args=args)
        def thread_api(empireMenu):

            try:
                start_restful_api(empireMenu=empireMenu, suppress=False, username=args.username, password=args.password, port=args.restport)
            except SystemExit as e:
                pass


        thread = helpers.KThread(target=thread_api, args=(main,))
        thread.daemon = True
        thread.start()
        sleep(2)

        def thread_websocket(empire_menu):
            try:
                start_sockets(empire_menu=empire_menu, port=int(args.socketport))
            except SystemExit as e:
                pass

        if args.notifications:
            thread2 = helpers.KThread(target=thread_websocket, args=(main,))
            thread2.daemon = True
            thread2.start()
            sleep(2)

        main.cmdloop()

    elif args.headless:
        # start an Empire instance and RESTful API and suppress output
        main = empire.MainMenu(args=args)

        try:
            start_restful_api(empireMenu=main, suppress=True, username=args.username, password=args.password, port=args.restport)
        except SystemExit as e:
            pass

    else:
        # normal execution
        main = empire.MainMenu(args=args)
        main.cmdloop()

    sys.exit()
